"""
结构代理 - 负责优化报告结构和章节安排
"""
import json
from typing import Dict, Any
from langchain.schema import HumanMessage
from .base_agent import BaseAgent
from .state import ReportState, AgentStatus, ReportStructure
from .utils import extract_json_from_text
import logging
import traceback

logger = logging.getLogger(__name__)

class StructureAgent(BaseAgent):
    """结构代理 - 负责优化报告结构和章节安排"""
    
    def __init__(self):
        super().__init__()
        self.name = "structure"
    
    def create_structure_prompt(self, plan_data: Dict[str, Any]) -> str:
        """创建结构优化提示词"""
        template = """
你是一个技术文档结构专家。请审查并优化以下报告结构，确保逻辑清晰、内容完整。

当前报告结构：
{plan_json}

请优化结构并返回改进后的JSON格式，要求：
1. 检查章节逻辑顺序是否合理
2. 确保内容覆盖完整，没有重要遗漏
3. 优化章节标题，使其更加专业和吸引人
4. 调整章节内容分配，确保平衡
5. 为需要代码的章节添加具体的代码要求

输出相同的JSON格式，但内容要更加完善和专业。
"""
        return self.create_prompt(template, plan_json=json.dumps(plan_data, ensure_ascii=False, indent=2))
    
    def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行结构化
        
        Args:
            state: 当前状态，包含 requirement 和 plan 等信息
            
        Returns:
            Dict[str, Any]: 结构化结果
        """
        try:
            logger.info("开始执行结构化...")
            logger.info(f"收到的状态数据: {state}")
            
            # 检查必需的字段
            if "requirement" not in state:
                error_msg = "状态中缺少 requirement 字段"
                logger.error(error_msg)
                return {"success": False, "error": error_msg, "status": "failed"}
            
            if "plan" not in state:
                error_msg = "状态中缺少 plan 字段"
                logger.error(error_msg)
                return {"success": False, "error": error_msg, "status": "failed"}
            
            requirement = state["requirement"]
            plan = state["plan"]
            logger.info(f"需求: {requirement}")
            logger.info(f"计划: {plan}")
            
            # 调用LLM生成结构
            messages = [
                {"role": "system", "content": "你是一个专业的文章结构设计专家。"},
                {"role": "user", "content": f"""
                根据以下信息设计文章结构：
                需求：{requirement}
                计划：{plan}
                """}
            ]
            
            logger.info("调用LLM生成结构...")
            result = self.call_llm(messages)
            logger.info(f"LLM返回结果: {result}")
            
            if result["success"]:
                logger.info("结构生成成功")
                return {
                    "success": True,
                    "content": result["content"],
                    "status": "completed"
                }
            else:
                logger.error(f"LLM调用失败: {result.get('error')}")
                return {
                    "success": False,
                    "error": result["error"],
                    "status": "failed"
                }
                
        except Exception as e:
            error_msg = f"结构化执行失败: {str(e)}"
            logger.error(error_msg)
            logger.error(f"错误详情: {traceback.format_exc()}")
            return {
                "success": False,
                "error": error_msg,
                "status": "failed"
            } 